//! PostgreSQL数据源实现
//! 
//! 基于sqlx的PostgreSQL数据源插件

use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use sqlx::{PgPool, Row, Column, TypeInfo};
use serde::{Serialize, Deserialize};
use serde_json::Value;

use crate::{DataSource, DataSourceQuery, DataSourceResult, DataSourceStats};

/// PostgreSQL数据源配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PostgreSqlConfig {
    /// 数据库URL
    pub database_url: String,
    /// 最大连接数
    pub max_connections: u32,
    /// 连接超时时间（秒）
    pub connect_timeout: u64,
    /// 查询超时时间（秒）
    pub query_timeout: u64,
}

impl Default for PostgreSqlConfig {
    fn default() -> Self {
        Self {
            database_url: "postgresql://localhost/test".to_string(),
            max_connections: 10,
            connect_timeout: 30,
            query_timeout: 60,
        }
    }
}

/// PostgreSQL数据源
pub struct PostgreSqlDataSource {
    config: PostgreSqlConfig,
    pool: Option<PgPool>,
}

impl PostgreSqlDataSource {
    /// 创建新的PostgreSQL数据源
    pub fn new(config: PostgreSqlConfig) -> Self {
        Self {
            config,
            pool: None,
        }
    }
    
    /// 获取连接池
    async fn get_pool(&self) -> anyhow::Result<&PgPool> {
        self.pool.as_ref().ok_or_else(|| {
            anyhow::anyhow!("数据源未连接")
        })
    }
    
    /// 将sqlx行转换为JSON值
    fn row_to_json(row: &sqlx::postgres::PgRow) -> anyhow::Result<HashMap<String, Value>> {
        let mut result = HashMap::new();
        
        for (i, column) in row.columns().iter().enumerate() {
            let column_name = column.name().to_string();
            let type_info = column.type_info();
            
            let value = match type_info.name() {
                "INT4" | "INTEGER" => {
                    match row.try_get::<Option<i32>, _>(i) {
                        Ok(Some(v)) => Value::Number(v.into()),
                        Ok(None) => Value::Null,
                        Err(_) => Value::Null,
                    }
                }
                "INT8" | "BIGINT" => {
                    match row.try_get::<Option<i64>, _>(i) {
                        Ok(Some(v)) => Value::Number(v.into()),
                        Ok(None) => Value::Null,
                        Err(_) => Value::Null,
                    }
                }
                "FLOAT4" | "REAL" => {
                    match row.try_get::<Option<f32>, _>(i) {
                        Ok(Some(v)) => Value::Number(serde_json::Number::from_f64(v as f64).unwrap_or_else(|| 0.into())),
                        Ok(None) => Value::Null,
                        Err(_) => Value::Null,
                    }
                }
                "FLOAT8" | "DOUBLE PRECISION" => {
                    match row.try_get::<Option<f64>, _>(i) {
                        Ok(Some(v)) => Value::Number(serde_json::Number::from_f64(v).unwrap_or_else(|| 0.into())),
                        Ok(None) => Value::Null,
                        Err(_) => Value::Null,
                    }
                }
                "BOOL" | "BOOLEAN" => {
                    match row.try_get::<Option<bool>, _>(i) {
                        Ok(Some(v)) => Value::Bool(v),
                        Ok(None) => Value::Null,
                        Err(_) => Value::Null,
                    }
                }
                "TEXT" | "VARCHAR" | "CHAR" => {
                    match row.try_get::<Option<String>, _>(i) {
                        Ok(Some(v)) => Value::String(v),
                        Ok(None) => Value::Null,
                        Err(_) => Value::Null,
                    }
                }
                _ => {
                    // 对于其他类型，尝试转换为字符串
                    match row.try_get::<Option<String>, _>(i) {
                        Ok(Some(v)) => Value::String(v),
                        Ok(None) => Value::Null,
                        Err(_) => Value::String(format!("未知类型: {}", type_info.name())),
                    }
                }
            };
            
            result.insert(column_name, value);
        }
        
        Ok(result)
    }
}

#[async_trait]
impl DataSource for PostgreSqlDataSource {
    fn name(&self) -> &str {
        "PostgreSQL"
    }

    fn description(&self) -> &str {
        "PostgreSQL数据库数据源"
    }

    async fn connect(&mut self) -> anyhow::Result<()> {
        let pool = PgPool::connect(&self.config.database_url)
            .await
            .map_err(|e| anyhow::anyhow!("连接PostgreSQL失败: {}", e))?;

        self.pool = Some(pool);
        Ok(())
    }
    
    async fn disconnect(&mut self) -> anyhow::Result<()> {
        if let Some(pool) = self.pool.take() {
            pool.close().await;
        }
        Ok(())
    }
    
    async fn is_connected(&self) -> bool {
        self.pool.is_some()
    }
    
    async fn execute_query(&mut self, query: DataSourceQuery) -> anyhow::Result<DataSourceResult> {
        let pool = self.get_pool().await?;
        
        let start_time = std::time::Instant::now();
        
        // 执行查询
        let rows = sqlx::query(&query.content)
            .fetch_all(pool)
            .await
            .map_err(|e| anyhow::anyhow!("查询执行失败: {}", e))?;
        
        let duration = start_time.elapsed();
        
        // 转换结果
        let mut data = Vec::new();
        for row in &rows {
            let json_row = Self::row_to_json(row)?;
            data.push(json_row);
        }
        
        Ok(DataSourceResult {
            success: true,
            data: Value::Array(data.into_iter().map(|row| {
                let map: serde_json::Map<String, Value> = row.into_iter().collect();
                Value::Object(map)
            }).collect()),
            error: None,
            stats: DataSourceStats {
                execution_time_ms: duration.as_millis() as f64,
                rows_processed: rows.len() as u64,
                bytes_processed: 0, // TODO: 计算实际字节数
            },
        })
    }

    async fn health_check(&self) -> anyhow::Result<Value> {
        let pool = self.get_pool().await?;

        match sqlx::query("SELECT 1 as status").fetch_one(pool).await {
            Ok(_) => Ok(serde_json::json!({
                "status": "healthy",
                "database": "PostgreSQL",
                "connected": true
            })),
            Err(e) => Ok(serde_json::json!({
                "status": "unhealthy",
                "database": "PostgreSQL",
                "connected": false,
                "error": e.to_string()
            })),
        }
    }
}

impl PostgreSqlDataSource {
    /// 获取数据库schema信息
    pub async fn get_schema(&self) -> anyhow::Result<Value> {
        let pool = self.get_pool().await?;
        
        // 查询所有表的信息
        let tables_query = r#"
            SELECT
                table_name,
                table_type
            FROM information_schema.tables
            WHERE table_schema = 'public'
            ORDER BY table_name
        "#;

        let table_rows = sqlx::query(tables_query)
            .fetch_all(pool)
            .await
            .map_err(|e| anyhow::anyhow!("获取表信息失败: {}", e))?;
        
        let mut tables = Vec::new();
        
        for table_row in table_rows {
            let table_name: String = table_row.get("table_name");
            let table_type: String = table_row.get("table_type");
            
            // 查询表的列信息
            let columns_query = r#"
                SELECT 
                    column_name,
                    data_type,
                    is_nullable,
                    column_default
                FROM information_schema.columns 
                WHERE table_schema = 'public' AND table_name = $1
                ORDER BY ordinal_position
            "#;
            
            let column_rows = sqlx::query(columns_query)
                .bind(&table_name)
                .fetch_all(pool)
                .await
                .map_err(|e| anyhow::anyhow!("获取列信息失败: {}", e))?;
            
            let mut columns = Vec::new();
            for column_row in column_rows {
                let column_name: String = column_row.get("column_name");
                let data_type: String = column_row.get("data_type");
                let is_nullable: String = column_row.get("is_nullable");
                let column_default: Option<String> = column_row.get("column_default");
                
                columns.push(serde_json::json!({
                    "name": column_name,
                    "type": data_type,
                    "nullable": is_nullable == "YES",
                    "default": column_default
                }));
            }
            
            tables.push(serde_json::json!({
                "name": table_name,
                "type": table_type,
                "columns": columns
            }));
        }
        
        Ok(serde_json::json!({
            "database_type": "PostgreSQL",
            "tables": tables
        }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_postgresql_data_source_creation() {
        let config = PostgreSqlConfig::default();
        let data_source = PostgreSqlDataSource::new(config);

        assert_eq!(data_source.name(), "PostgreSQL");
        assert_eq!(data_source.description(), "PostgreSQL数据库数据源");
        assert!(!data_source.is_connected().await);
    }

    #[tokio::test]
    async fn test_postgresql_config_default() {
        let config = PostgreSqlConfig::default();
        assert_eq!(config.database_url, "postgresql://localhost/test");
        assert_eq!(config.max_connections, 10);
        assert_eq!(config.connect_timeout, 30);
        assert_eq!(config.query_timeout, 60);
    }

    #[tokio::test]
    async fn test_row_to_json() {
        // 这个测试需要实际的数据库连接，所以我们只测试基本功能
        // 在实际环境中，可以使用测试数据库进行完整测试
        
        // 测试基本的数据源创建和配置
        let config = PostgreSqlConfig {
            database_url: "postgresql://test:test@localhost/testdb".to_string(),
            max_connections: 5,
            connect_timeout: 10,
            query_timeout: 30,
        };
        
        let data_source = PostgreSqlDataSource::new(config.clone());
        assert_eq!(data_source.config.database_url, config.database_url);
        assert_eq!(data_source.config.max_connections, 5);
    }
}
